package http.secure.common.util.steganography;

import lombok.extern.slf4j.Slf4j;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.nio.charset.StandardCharsets;

@Slf4j
public class ImageSteganography {
    /**
     * 将数据隐藏到图片中
     *
     * @param imageBytes 原始图片数据
     * @param message    要隐藏的数据
     * @return 隐写后的图片二进制数据
     */
    public static byte[] hideMessage(byte[] imageBytes, String message) {
        try {
            // 读取原始图片
            BufferedImage image = ImageIO.read(new ByteArrayInputStream(imageBytes));
            int width = image.getWidth();
            int height = image.getHeight();
            // 将数据转换为二进制
            byte[] dataBytes = message.getBytes(StandardCharsets.UTF_8);
            String binaryData = bytesToBinary(dataBytes);
            // 在数据前添加长度信息（32位表示数据长度）
            String binaryString = Integer.toBinaryString(dataBytes.length);
            String lengthBinary = String.format("%32s", binaryString).replace(' ', '0');
            binaryData = lengthBinary + binaryData;
            // 检查图片容量是否足够
            if (binaryData.length() > width * height * 3) {
                throw new IllegalStateException("错误：图片容量不足以隐藏数据");
            }
            // 隐藏数据
            hideBinaryData(image, binaryData);
            // 保存新图片
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            ImageIO.write(image, "png", bos);
            return bos.toByteArray();
        } catch (Exception e) {
            throw new IllegalStateException("图片隐写失败：" + e.getMessage());
        }
    }

    /**
     * 将字节数组转换为二进制字符串
     */
    private static String bytesToBinary(byte[] bytes) {
        StringBuilder sb = new StringBuilder();
        for (byte b : bytes) {
            for (int i = 7; i >= 0; i--) {
                sb.append((b >> i) & 1);
            }
        }
        return sb.toString();
    }

    /**
     * 将二进制数据隐藏到图片中
     */
    private static void hideBinaryData(BufferedImage image, String binaryData) {
        int width = image.getWidth();
        int height = image.getHeight();
        int dataIndex = 0;
        outerLoop:
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                int rgb = image.getRGB(x, y);
                int alpha = (rgb >> 24) & 0xFF;
                int red = (rgb >> 16) & 0xFF;
                int green = (rgb >> 8) & 0xFF;
                int blue = rgb & 0xFF;
                // 修改红色通道
                if (dataIndex < binaryData.length()) {
                    red = modifyLSB(red, binaryData.charAt(dataIndex) - '0');
                    dataIndex++;
                }
                // 修改绿色通道
                if (dataIndex < binaryData.length()) {
                    green = modifyLSB(green, binaryData.charAt(dataIndex) - '0');
                    dataIndex++;
                }
                // 修改蓝色通道
                if (dataIndex < binaryData.length()) {
                    blue = modifyLSB(blue, binaryData.charAt(dataIndex) - '0');
                    dataIndex++;
                }
                // 设置修改后的像素值
                int newRGB = (alpha << 24) | (red << 16) | (green << 8) | blue;
                image.setRGB(x, y, newRGB);
                // 如果数据已经全部隐藏完毕，跳出循环
                if (dataIndex >= binaryData.length()) {
                    break outerLoop;
                }
            }
        }
    }

    /**
     * 修改数字的最低有效位
     */
    private static int modifyLSB(int value, int bit) {
        // 清除最低位，然后设置为新的值
        return (value & 0xFE) | bit;
    }
}